# 例3-1
import numpy as np

# 模拟向量数据
np.random.seed(42)
data_vectors = np.random.rand(100, 5)  # 100个5维向量


# -------------------------------
# 倒排索引实现
# -------------------------------
class InvertedIndex:
    def __init__(self):
        self.index = {}

    def add_to_index(self, vector_id, bucket):
        """将向量ID添加到倒排索引的对应桶"""
        if bucket not in self.index:
            self.index[bucket] = []
        self.index[bucket].append(vector_id)

    def search(self, bucket):
        """检索某个桶中的所有向量ID"""
        return self.index.get(bucket, [])


# 构建倒排索引
inverted_index = InvertedIndex()
print("data_vectors: ", data_vectors[:2])

for i, vector in enumerate(data_vectors):
    bucket = int(vector[0] * 10)  # 简单分桶：根据第一个维度值划分
    inverted_index.add_to_index(i, bucket)

# 倒排索引检索
bucket_to_search = 2
result_ids = inverted_index.search(bucket_to_search)
print(f"倒排索引中桶 {bucket_to_search} 的向量ID:", result_ids)
print("bucket nums: ", inverted_index.index.keys())
